package com.demo.function.shardingsphere.config;

import com.google.common.collect.Lists;
import lombok.extern.slf4j.Slf4j;
import org.apache.shardingsphere.core.rule.BindingTableRule;
import org.apache.shardingsphere.core.rule.TableRule;
import org.apache.shardingsphere.shardingjdbc.jdbc.core.connection.ShardingConnection;
import org.apache.shardingsphere.shardingjdbc.jdbc.core.context.ShardingRuntimeContext;
import org.apache.shardingsphere.shardingjdbc.jdbc.core.datasource.ShardingDataSource;
import org.apache.shardingsphere.shardingjdbc.spring.boot.SpringBootConfiguration;
import org.springframework.boot.autoconfigure.AutoConfigureAfter;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.context.EnvironmentAware;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.env.Environment;
import org.springframework.lang.NonNull;

import javax.sql.DataSource;
import java.sql.ResultSet;
import java.util.List;
import java.util.Set;

/**
 * 绑定表规则和表对应的数据源
 * @author K
 */
@Slf4j
@ConditionalOnProperty(name = {"sharding.binding-tables-enable"}, havingValue = "true")
@Configuration
@AutoConfigureAfter(SpringBootConfiguration.class)
public class ShardingTablesAware implements EnvironmentAware {

    private final static String GET_TABLE_SQL = "select table_name from information_schema.tables where table_schema='%s'";
    private final static String TABLE_NAME_COLUMN = "table_name";

    private final DataSource dataSource;

    public ShardingTablesAware(DataSource dataSource) {
        this.dataSource = dataSource;
    }

    @Override
    public void setEnvironment(@NonNull Environment environment) {
        log.info("加载binding-table和table-rule");
        ShardingDataSource shardingDataSource = (ShardingDataSource) dataSource;
        ShardingRuntimeContext runtimeContext = shardingDataSource.getRuntimeContext();
        Set<String> dsSet = shardingDataSource.getDataSourceMap().keySet();
        for (String ds : dsSet) {
            String schema = runtimeContext.getMetaData().getDataSources().getDataSourceMetaData(ds).getCatalog();
            try (ShardingConnection connection = shardingDataSource.getConnection();
                 ResultSet resultSet = connection.createStatement().executeQuery(String.format(GET_TABLE_SQL, schema))) {

                List<String> tableNameList = Lists.newArrayList();
                List<TableRule> tableRules = Lists.newArrayList();
                while (resultSet.next()) {
                    String tableName = resultSet.getString(TABLE_NAME_COLUMN);
                    tableNameList.add(tableName);
                }
                for (String tableName : tableNameList) {
                    if (shardingDataSource.getRuntimeContext().getRule().findTableRuleByActualTable(tableName).isEmpty()) {
                        TableRule rule = new TableRule(dsSet, tableName);
                        shardingDataSource.getRuntimeContext().getRule().getTableRules().add(rule);

                        tableRules.add(rule);
                    }
                }

                shardingDataSource.getRuntimeContext().getRule().getBindingTableRules().add(new BindingTableRule(tableRules));

            } catch (Exception e) {
                log.error("shardingshpere 加载异常", e);
            }
        }
    }
}
